import path from 'path';
import { Injectable } from '@nestjs/common';
import { CheerioWebBaseLoader } from '@langchain/community/document_loaders/web/cheerio';
import { JSONLoader } from 'langchain/document_loaders/fs/json';
import { TextLoader } from 'langchain/document_loaders/fs/text';
import { CSVLoader } from '@langchain/community/document_loaders/fs/csv';
import { DocxLoader } from '@langchain/community/document_loaders/fs/docx';
import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf';

import { CreateDocumentDto } from './dto/create-document.dto';
import { WeaviateService } from 'src/core/weaviate/weaviate.service';
import { StoreService } from '../store/store.service';

@Injectable()
export class DocumentService {
  constructor(
    private readonly weaviateService: WeaviateService,
    private readonly storeService: StoreService,
  ) {}

  async create(createDocumentDto: CreateDocumentDto) {
    const { file } = createDocumentDto;

    const store = await this.storeService.findOne(file);

    const filePath = path.join(process.cwd(), 'static/store', `${store.id}${store.type}`);

    const loader = this.getLoader(store.type, filePath);

    const docs = await loader.load();

    for (let i = 0; i < docs.length; i++) {
      const element = docs[i];
      element.metadata = { ...element.metadata, source: file };
    }

    const ids = await this.weaviateService.addDocuments(docs);

    return ids;
  }

  async load(createDocumentDto: CreateDocumentDto) {
    const { file } = createDocumentDto;

    const laoder = new CheerioWebBaseLoader(file);

    const docs = await laoder.load();

    await this.weaviateService.addDocuments(docs);

    return '成功加载文档并转换为向量存储';
  }

  async findAll(current: number = 1, pageSize: number = 10, source: string) {
    return this.weaviateService.findAll(current, pageSize, source);
  }

  async removeAll() {
    await this.weaviateService.deleteDocuments();

    return '删除成功';
  }

  getLoader(ext: string, filePath: string) {
    switch (ext) {
      case '.txt':
      case '.md':
        return new TextLoader(filePath);
      case '.pdf':
        return new PDFLoader(filePath);
      case '.csv':
        return new CSVLoader(filePath);
      case '.json':
        return new JSONLoader(filePath);
      case '.docx':
        return new DocxLoader(filePath);
      case '.doc':
        return new DocxLoader(filePath, { type: 'doc' });
      default:
        throw new Error(`不支持的文件类型: ${ext}`);
    }
  }
}
